import onnxruntime
from PIL import Image,ImageDraw
from torchvision import transforms
import torch
import numpy as np
import cv2
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords,
    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)
# ort_session = onnxruntime.InferenceSession("torch_model.onnx")
ort_session = onnxruntime.InferenceSession("yolov5s.onnx")
print("Exported model has been tested with ONNXRuntime, and the result looks good!")

img_sr = Image.open(r"../inference/images/zidane.jpg")
img_sr = img_sr.resize((640, 640),Image.ANTIALIAS)
img_numpy = cv2.cvtColor(np.array(img_sr), cv2.COLOR_RGB2BGR)
# img_numpy = img_numpy[:,:,::-1].transpose(2, 0, 1)
img_numpy = np.rollaxis(img_numpy, 2)
# device = select_device(opt.device)
img = torch.from_numpy(img_numpy).to(torch.device('cuda:0'))
# img = img.half()  # uint8 to fp16/32
img = img.float()
img /= 255.0  # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
    img = img.unsqueeze(0)
# tf = transforms.Compose([
#     transforms.Resize((512,640)),
#     transforms.ToTensor()
# ])
# img_tensor = tf(img)
# ort_inputs = {ort_session.get_inputs()[0].name: img_tensor[None].numpy()}


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}
out = ort_session.run(None, ort_inputs)
pred_0 = out[0].reshape([-1, 85])
pred_1 = out[1].reshape([-1, 85])
pred_2 = out[2].reshape([-1, 85])
cat_pred = np.concatenate((pred_0, pred_1), axis=0)
cat_pred = np.concatenate((cat_pred, pred_2), axis=0)
cat_pred_t = torch.FloatTensor(cat_pred)
cat_pred_t = cat_pred_t.unsqueeze(0)
# pred = torch.tensor(ort_session.run(None, ort_inputs)[0])
det = non_max_suppression(cat_pred_t, 0.4, 0.5, classes=None, agnostic=False)
# img = img.resize((640, 512))
# Imgdraw = ImageDraw.Draw(img)
Imgdraw = ImageDraw.Draw(img_sr)

for box in det[0]:
    b = box.cpu().detach().long().numpy()
    print(b)
    Imgdraw.rectangle((b[0],b[1],b[2],b[3]))

# img.show()
img_sr.show()
